load("//bazel:tachyon.bzl", "if_aarch64", "if_has_avx512", "if_x86_64")
load("//bazel:tachyon_cc.bzl", "tachyon_avx512_defines", "tachyon_cc_library")
load("//tachyon/math/finite_fields/generator/ext_field_generator:build_defs.bzl", "generate_fp4s")
load("//tachyon/math/finite_fields/generator/prime_field_generator:build_defs.bzl", "SUBGROUP_GENERATOR", "generate_fft_prime_fields")

package(default_visibility = ["//visibility:public"])

generate_fp4s(
    name = "baby_bear4",
    base_field = "BabyBear",
    base_field_degree = 1,
    base_field_hdr = "tachyon/math/finite_fields/baby_bear/internal/baby_bear.h",
    class_name = "BabyBear4",
    is_packed = False,
    namespace = "tachyon::math",
    # See https://github.com/Plonky3/Plonky3/blob/d9ef390/baby-bear/src/baby_bear.rs#L80.
    non_residue = ["11"],
    deps = [":baby_bear"],
)

generate_fp4s(
    name = "packed_baby_bear4",
    base_field = "PackedBabyBear",
    base_field_degree = 1,
    base_field_hdr = "tachyon/math/finite_fields/baby_bear/internal/packed_baby_bear.h",
    class_name = "PackedBabyBear4",
    is_packed = True,
    namespace = "tachyon::math",
    # See https://github.com/Plonky3/Plonky3/blob/d9ef390/baby-bear/src/baby_bear.rs#L80.
    non_residue = ["11"],
    deps = [":packed_baby_bear"],
)

generate_fft_prime_fields(
    name = "baby_bear",
    class_name = "BabyBear",
    # 2³¹ - 2²⁷ + 1
    # Hex: 0x78000001
    modulus = "2013265921",
    namespace = "tachyon::math",
    subgroup_generator = "//tachyon/math/finite_fields/baby_bear:" + SUBGROUP_GENERATOR,
    use_montgomery = True,
)

tachyon_cc_library(
    name = "packed_baby_bear",
    hdrs = ["packed_baby_bear.h"],
    defines = tachyon_avx512_defines(),
    deps = [
        ":packed_baby_bear_neon",
        "//tachyon/build:build_config",
        "//tachyon/math/finite_fields:extended_packed_field_traits_forward",
        "//tachyon/math/finite_fields:finite_field_traits",
        "//tachyon/math/matrix:prime_field_num_traits",
    ] + if_has_avx512(
        [":packed_baby_bear_avx512"],
        [":packed_baby_bear_avx2"],
    ),
)

tachyon_cc_library(
    name = "packed_baby_bear_avx2",
    srcs = if_x86_64(["packed_baby_bear_avx2.cc"]),
    hdrs = if_x86_64(["packed_baby_bear_avx2.h"]),
    copts = if_x86_64(["-mavx2"]),
    deps = [
        ":baby_bear",
        "//tachyon:export",
        "//tachyon/math/finite_fields:packed_prime_field32_avx2",
        "//tachyon/math/finite_fields:packed_prime_field_base",
    ],
)

tachyon_cc_library(
    name = "packed_baby_bear_avx512",
    srcs = if_x86_64(["packed_baby_bear_avx512.cc"]),
    hdrs = if_x86_64(["packed_baby_bear_avx512.h"]),
    copts = if_x86_64([
        "-mavx512f",
        # NOTE(chokobole): See https://gitlab.com/libeigen/eigen/-/blob/0b51f76/Eigen/src/Core/util/ConfigureVectorization.h#L248.
        "-mfma",
    ]),
    deps = [
        ":baby_bear",
        "//tachyon:export",
        "//tachyon/math/finite_fields:packed_prime_field32_avx512",
        "//tachyon/math/finite_fields:packed_prime_field_base",
    ],
)

tachyon_cc_library(
    name = "packed_baby_bear_neon",
    srcs = if_aarch64(["packed_baby_bear_neon.cc"]),
    hdrs = if_aarch64(["packed_baby_bear_neon.h"]),
    deps = [
        ":baby_bear",
        "//tachyon:export",
        "//tachyon/math/finite_fields:packed_prime_field32_neon",
        "//tachyon/math/finite_fields:packed_prime_field_base",
    ],
)
